package yz.cp.back.websocket.config;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.messaging.simp.config.ChannelRegistration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;
import org.springframework.web.socket.server.support.DefaultHandshakeHandler;
import yz.cp.back.websocket.handler.AuthHandshakeInterceptor;
import yz.cp.back.websocket.handler.ChannelInInterceptor;

import java.security.Principal;
import java.util.Map;

/**
 * Created by xukangkang on 2019/3/14.
 */
@Configuration
@EnableWebSocketMessageBroker
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {
    private static long HEART_BEAT=5000;
    @Override
    public void configureMessageBroker(MessageBrokerRegistry config) {
        ThreadPoolTaskScheduler te = new ThreadPoolTaskScheduler();
        te.setPoolSize(10);
        te.setThreadNamePrefix("wss-heartbeat-thread-");
        te.initialize();
        config.enableSimpleBroker("/topic","/queue")
        .setHeartbeatValue(new long[]{HEART_BEAT,HEART_BEAT}).setTaskScheduler(te);
        //定义了服务端接收地址的前缀，也即客户端给服务端发消息的地址前缀
        config.setApplicationDestinationPrefixes("/app"); //应用请求前缀
        //config.setUserDestinationPrefix("/ricky/");//推送用户前缀
    }

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {
        registry.addEndpoint("/ws/sockjs")
                .setAllowedOrigins("*")
                .addInterceptors(authHandshakeInterceptor())
                .setHandshakeHandler(new DefaultHandshakeHandler(){
                    @Override
                    protected Principal determineUser(ServerHttpRequest request, WebSocketHandler wsHandler, Map<String, Object> attributes) {
                        //设置认证用户
                        return (Principal)attributes.get("user");
                    }
                })
               .withSockJS();
        registry.addEndpoint("/ws")
                .setAllowedOrigins("*")
                .addInterceptors(authHandshakeInterceptor())
                .setHandshakeHandler(new DefaultHandshakeHandler(){
                    @Override
                    protected Principal determineUser(ServerHttpRequest request, WebSocketHandler wsHandler, Map<String, Object> attributes) {
                        //设置认证用户
                        return (Principal)attributes.get("user");
                    }
                });
    }

    @Override
    public void configureClientInboundChannel(ChannelRegistration registration) {
        registration.interceptors(channelInInterceptor());
    }

    @Bean
    public ChannelInInterceptor channelInInterceptor(){
        return new ChannelInInterceptor();
    }
    @Bean
    public AuthHandshakeInterceptor authHandshakeInterceptor(){
        return new AuthHandshakeInterceptor();
    }

}